package cn.sskxyz.mongodb.multi;

import java.io.InputStream;

import org.bson.Document;
import org.springframework.data.mongodb.core.convert.MongoConverter;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.data.mongodb.gridfs.GridFsObject;
import org.springframework.data.mongodb.gridfs.GridFsResource;
import org.springframework.data.mongodb.gridfs.GridFsTemplate;
import org.springframework.data.mongodb.util.BsonUtils;
import org.springframework.util.Assert;

import com.mongodb.client.MongoDatabase;
import com.mongodb.client.gridfs.GridFSBucket;
import com.mongodb.client.gridfs.GridFSBuckets;
import com.mongodb.client.gridfs.GridFSFindIterable;
import com.mongodb.client.gridfs.model.GridFSFile;
import com.mongodb.client.gridfs.model.GridFSUploadOptions;

public class MultiGridFsTemplate extends GridFsTemplate {

	private MultiGridFsDatabaseFactory gridFsDatabaseFactory;

	public MultiGridFsTemplate(MultiGridFsDatabaseFactory gridFsDatabaseFactory, MongoConverter converter) {
		super(gridFsDatabaseFactory, converter);
		this.gridFsDatabaseFactory = gridFsDatabaseFactory;
	}

	@Override
	public GridFSFindIterable find(Query query) {
		Assert.notNull(query, "Query must not be null!");

		Document queryObject = getMappedQuery(query.getQueryObject());
		Document sortObject = getMappedQuery(query.getSortObject());

		GridFSFindIterable iterable = getGridFs().find(queryObject).sort(sortObject);

		if (query.getSkip() > 0) {
			iterable = iterable.skip(Math.toIntExact(query.getSkip()));
		}

		if (query.getLimit() > 0) {
			iterable = iterable.limit(query.getLimit());
		}

		return iterable;
	}

	@Override
	public void delete(Query query) {
		for (GridFSFile gridFSFile : find(query)) {
			getGridFs().delete(gridFSFile.getId());
		}
	}

	@Override
	public GridFsResource getResource(GridFSFile file) {
		Assert.notNull(file, "GridFSFile must not be null!");

		return new GridFsResource(file, getGridFs().openDownloadStream(file.getId()));
	}

	@Override
	public <T> T store(GridFsObject<T, InputStream> upload) {
		String contentType = (String) upload.getOptions().getMetadata().get("_contentType");
		GridFSUploadOptions uploadOptions = computeUploadOptionsFor(contentType, upload.getOptions().getMetadata());

		if (upload.getOptions().getChunkSize() > 0) {
			uploadOptions.chunkSizeBytes(upload.getOptions().getChunkSize());
		}

		if (upload.getFileId() == null) {
			return (T) getGridFs().uploadFromStream(upload.getFilename(), upload.getContent(), uploadOptions);
		}

		getGridFs().uploadFromStream(BsonUtils.simpleToBsonValue(upload.getFileId()), upload.getFilename(),
				upload.getContent(), uploadOptions);
		return upload.getFileId();
	}

	private GridFSBucket getGridFs() {
		MongoDatabase db = gridFsDatabaseFactory.getMongoDatabase();
		String bucket = gridFsDatabaseFactory.getMongoProperties().getGridfs().getBucket();
		return bucket == null ? GridFSBuckets.create(db) : GridFSBuckets.create(db, bucket);
	}

}
